# python imports
import math
from time import time
from random import Random

try:
  import numpy
except:
  # its OK if numpy is not installed
  pass

# local imports
from misc.util import Counter, DefaultDict

# pyblog imports
import glob
from glob import GROUPBY_NOTHING, GROUPBY_UNUSED_ARGS, GROUPBY_INDICATOR_FN
from world import World
from proposer import flip_var

QUERY_DISTRIBUTION, QUERY_LAST, QUERY_SEQUENCE = range(3)

def query(querylist, obslist, init_vals=None, burnin=100, scans=1000,
          seed=None, trace=False, stats=True, outtyp=QUERY_DISTRIBUTION,
          debug=True, deterministic=False, statsobj=None):
  """
  Use sampling methods to estimate a distribution for each expression
  in the list of queries. In one scan, all the variables in the model
  are sampled (flipped) at least once. At the end of each scan,
  a sample of the variables to be queried are taken for computing their
  distribution.

  querylist   - the list of queries [RandomVar(..), ..]
  obslist     - the list of observations [RandomVar(..) == val, ..]
  init_vals   - the list of initial values [RandomVar(..) == val, ..]
  burnin      - number of initial samples to discard
  scans       - number of times to sample each variable in the model
  seed        - the random number seed (defaults to current time)
  trace       - print some tracing in each world (False)
  stats       - print stats on the performance (True)
  outtyp      - QUERY_DISTRIBUTION : return the distribution
                QUERY_LAST         : return the last value
                QUERY_SEQUENCE     : return the sequence of values
  debug       - debug on exception (True)
  deterministic - Remove non-determinism e.g. when converting a set
                  datatype to a list, sort the elements of the set etc. (False)
  statsobj    - if a dictionary is passed in, it will be filled with stats
  """
  # we will put the entire code in one big try .. finally block to clean
  # the curr_world global on abnormal exit
  try:
    return _query(querylist, obslist, init_vals, burnin, scans,
                  seed, trace, stats, outtyp, debug, deterministic, statsobj)
  finally:
    # clear, the current world global for future use
    glob.curr_world = None
  
def _query(querylist, obslist, init_vals, burnin, scans,
           seed, trace, stats, outtyp, debug, deterministic, user_statsobj):
  try:
    # create a random number generator
    if seed is None:
      seed = int(time())
      if stats:
        print "Seed:", seed
    rand = Random(seed)
    # set the numpy seed but its OK if numpy doesn't exist
    try:
      numpy.random.seed(seed)
    except:
      pass

    if outtyp == QUERY_DISTRIBUTION:
      # initialize a Counter object to track outcomes for each queried variable
      outcome = [Counter() for _ in querylist]
    elif outtyp == QUERY_SEQUENCE:
      outcome = [[] for _ in querylist]
    
    t1 = time()
    # create a new world
    glob.curr_world = World(rand, trace, deterministic)
    init_time = time() - t1

    # set the initial values
    if init_vals is not None:
      glob.curr_world.set_init_vals(init_vals)
    
    # set the evidence
    if obslist:
      glob.curr_world.set_observation(obslist)

    # initialize the scan counters
    glob.curr_world.init_scan()
    
    statsobj = DefaultDict(default=lambda x: Counter())
    statsobj["init-time"] = init_time

    if stats:
      print "Initial World took %.1f seconds" % init_time,
      print " -- %d variables, %d sampleable" \
            % (len(glob.curr_world.variables),
               len(glob.curr_world.sampleable_vars))

    
    # iterate
    reject_count = flip_count = 0
    t1 = time()
    # keep iterating until the desired number of scans
    while glob.curr_world.scan_cnt < scans:
      
      # pick the next variable
      var = glob.curr_world.next_scan_var()
      
      # flip the next variable if there is one
      if var is not None:
        reject_count += flip_var(var, glob.curr_world, rand, trace, statsobj)
        flip_count += 1
        continue
      
      # otherwise, a scan has ended we need to collect a sample
      
      # discard the initial burnin samples
      if glob.curr_world.scan_cnt < burnin:
        pass
      
      elif outtyp in [QUERY_DISTRIBUTION, QUERY_SEQUENCE]:
        # evaluate the queries in this world
        curr_outcomes = glob.curr_world.query_and_forget(querylist)
        if trace:
          print "[%d]: OUTCOMES" % glob.curr_world.scan_cnt, str(curr_outcomes)
        # update the statistics of the queries
        if outtyp == QUERY_DISTRIBUTION:
          for qnum, o in enumerate(curr_outcomes):
            outcome[qnum][o] += 1
        elif outtyp == QUERY_SEQUENCE:
          for qnum, o in enumerate(curr_outcomes):
            outcome[qnum].append(o)
      
      glob.curr_world.next_scan()       # start a new scan            

    t2 = time()
    scan_time = t2-t1
    statsobj["scan-time"] = scan_time
    
    if stats and flip_count:
      print "%d scans in %d s, %.1f flips per sec, %d "\
            "%% flips rejected" \
            % (glob.curr_world.scan_cnt, int(scan_time),
               float(flip_count)/(scan_time+1e-3),
               100 * reject_count / flip_count)
      for fn, fnstats in statsobj.iteritems():
        if callable(fn):
          t0, t1, t2 = fnstats[0], fnstats[1], fnstats[2]
          del fnstats[0], fnstats[1], fnstats[2]
          print fn.__name__, fnstats,
          print "avg %.2e" % ( t1/t0), \
                "sd %.2e" % ( math.sqrt(t2/t0 - (t1/t0)**2))

      for fn, vardef in glob.var_def.iteritems():
        if vardef.groupby == GROUPBY_UNUSED_ARGS:
          print fn.__name__, \
                ": GBY-UNUSED %d args" % (vardef.argcount - vardef.used_args)
        elif vardef.groupby == GROUPBY_INDICATOR_FN:
          print fn.__name__, ": GBY-IND %s" % (vardef.ind_fn.__name__),
          if vardef.bound_fns is not None:
            print "bound fns:", str([bndfn.__name__ \
                                     for bndfn in vardef.bound_fns])
          else:
            print


    t1 = time()
    
    # finish up with generating the answer
    if outtyp == QUERY_LAST:
      outcome = glob.curr_world.query_and_forget(querylist)

    elif outtyp == QUERY_DISTRIBUTION:
      # normalize and return the outcomes as a distribution
      for o in outcome:
        o.normalize()

    elif outtyp == QUERY_SEQUENCE:
      pass

    else:
      raise TypeError("Unknown query output type %d" % outtyp)

    final_time = time() - t1
    statsobj["final-time"] = final_time
    
    if stats:
      print "Final Query took: %.1f seconds" % (final_time)
      
    # at the end of the iterations, clear the global variable curr_world
    glob.curr_world = None

    # copy the stats to the user stats
    if user_statsobj is not None:
      for key, val in statsobj.iteritems():
        user_statsobj[key] = val
  
    return outcome

  except:
    if debug:
      import pdb, traceback, sys
      print
      print "-" * 60
      print "Exception inside pyblog"
      print "-" * 60
      print
      traceback.print_exc(file=sys.stdout)
      print
      print "-" * 60
      print "Please debug or type 'q' to exit"
      print "-" * 60
      print
      pdb.post_mortem(sys.exc_traceback)
    raise

def configure(param, value):
  assert(glob.curr_world is None)       # can't be inside a world

  if not hasattr(glob, param):
    raise ValueError("unknown configuration parameter:" + str(param))

  #print "configure:", param, ":", getattr(glob, param), "-->", value
  
  setattr(glob, param, value)

  # clear out vardef to ensure that the new parameter is picked up
  for vardef in glob.var_def.itervalues():
    vardef.groupby = None
  
  
